package main

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/axgle/mahonia"
	"github.com/cakturk/go-netstat/netstat"
	"github.com/kardianos/service"
	"github.com/oschwald/geoip2-golang"
)

type program struct{}

func (p *program) Start(s service.Service) error {
	go p.run()
	return nil
}

// 需要防护的端口
var ports = []uint16{}

// 读取配置文件把需要防护的端口放到ports里面
func getPorts() {
	root_path := GetRootPath()
	config_path := path.Join(root_path, "ports.txt")
	if _, err := os.Stat(config_path); err == nil {
		// 读取配置文件
		f, err := os.Open(config_path)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		contents, err := os.ReadFile(config_path)
		if err != nil {
			log.Fatal(err)
		}
		for _, port := range strings.Split(string(contents), ",") {
			portNumber, err := strconv.Atoi(port)
			if err == nil && portNumber > 0 && portNumber < 65535 {
				ports = append(ports, uint16(portNumber))
			}
		}
	}
	log.Println("防护端口列表：", ports)

}
func inPorts(port uint16) bool {
	for _, p := range ports {
		if p == port {
			return true
		}
	}
	return false
}

func (p *program) run() {
	// 程序主进程
	// 创建一个字符串型的管道记录ip
	ipChan := make(chan string, 100)
	//10s 检测一次当前的网络连接
	ticker := time.NewTicker(10 * time.Second)
	go func() {
		for range ticker.C {
			s, _ := DisplaySocks(8899)
			for _, e := range s {
				remoteip := e.RemoteAddr.IP.String()
				c := GeoIp(remoteip)
				if c != "CN" && c != "" {
					log.Printf("%v:%v待封锁\n", remoteip, c)
					ipChan <- remoteip
				}
			}
		}
	}()

	// 把管道中的ip拉入黑名单
	for {
		// 接收数据
		out, ok := <-ipChan
		// 判断通道是否关闭,如果通道关闭，则ok为false
		if !ok {
			fmt.Println("通道已关闭")
			break
		}
		BlockIP(out)
		log.Printf("已封锁%s\n", out)
	}
}

func (p *program) Stop(s service.Service) error {
	s.Stop()
	return nil
}

func init() {
	// 日志
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
	logPath := setLogPath()
	logFile, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		log.Panic("打开日志文件异常")
	}
	log.SetOutput(logFile)
	// 端口列表读取
	getPorts()
}

func main() {
	svConfig := &service.Config{
		Name:        "Suda-Guard",
		DisplayName: "速达守护服务",
		Description: "让速达更安全地运行",
	}
	prg := &program{}
	s, err := service.New(prg, svConfig)
	if err != nil {
		log.Fatal(err)
	}
	if len(os.Args) > 1 {
		if os.Args[1] == "install" {
			s.Install()
			log.Println("服务安装成功")
			return
		}
		if os.Args[1] == "remove" {
			s.Uninstall()
			log.Println("服务卸载成功")
			return
		}
	}
	if err = s.Run(); err != nil {
		log.Fatal(err)
	}
}

func GetRootPath() string {
	exePath, err := os.Executable()
	if err != nil {
		log.Fatal(err)
	}
	res, _ := filepath.EvalSymlinks(filepath.Dir(exePath))
	return res

}

func GetDateString() string {
	now := time.Now()

	return now.Format("2006-01-02")
}

func setLogPath() string {
	root_path := GetRootPath()
	date_now := GetDateString()

	return path.Join(root_path, date_now+".log")

}
func getGeoPath() string {
	root_path := GetRootPath()

	return path.Join(root_path, "GeoLite2-City.mmdb")
}
func DisplaySocks(port uint16) ([]netstat.SockTabEntry, error) {
	tabs, err := netstat.TCPSocks(func(s *netstat.SockTabEntry) bool {
		// return s.LocalAddr.Port == port
		return inPorts(s.LocalAddr.Port)
		// return true
	})

	return tabs, err
}

func GeoIp(ip string) string {
	db, err := geoip2.Open(getGeoPath())
	if err != nil {
		log.Panic(err)
	}
	defer db.Close()
	// If you are using strings that may be invalid, check that ip is not nil
	ip2 := net.ParseIP(ip)
	record, err := db.City(ip2)
	// record, err := db.Country(ip)
	if err != nil {
		log.Panic(err)
	}
	// fmt.Printf("Portuguese (BR) city name: %v\n", record.City.Names["pt-BR"])
	// fmt.Printf("English subdivision name: %v\n", record.Subdivisions[0].Names["en"])
	// fmt.Printf("Russian country name: %v\n", record.Country.Names["ru"])
	// fmt.Printf("ISO country code: %v\n", record.Country.IsoCode)
	// fmt.Printf("Time zone: %v\n", record.Location.TimeZone)
	// fmt.Printf("Coordinates: %v, %v\n", record.Location.Latitude, record.Location.Longitude)
	// fmt.Printf("国家isocode:%v\n", record.Country.IsoCode)
	// fmt.Printf("城市:%v\n", record.City.Names)

	return record.Country.IsoCode
}

func BlockIP(ip string) {
	if notHasRule(ip) {
		name := "block" + ip
		shell := "netsh advfirewall firewall add rule name=%s dir=in action=block protocol=tcp remoteip=%s"
		cmdstr := fmt.Sprintf(shell, name, ip)
		cmd(cmdstr)
	}
}

func notHasRule(ip string) bool {
	name := "block" + ip
	shell := "netsh advfirewall firewall show rule name=%s"
	cmdstr := fmt.Sprintf(shell, name)

	rt := cmd(cmdstr)

	return strings.Contains(rt, "没有与指定标准相匹配的规则") || strings.Contains(rt, "No rules match")
}

func cmd(bat string) string {
	cmd := exec.Command("cmd")
	in := bytes.NewBuffer(nil)
	cmd.Stdin = in //绑定输入
	var out bytes.Buffer
	cmd.Stdout = &out   //绑定输出
	in.WriteString(bat) //写入你的命令，可以有多行，"\n"表示回车
	in.WriteString("\n")
	err := cmd.Start()
	if err != nil {
		log.Fatal(err)
	}
	// log.Println(cmd.Args)
	err = cmd.Wait()
	if err != nil {
		log.Printf("Command finished with error: %v", err)
	}
	rt := mahonia.NewDecoder("gbk").ConvertString(out.String())
	fmt.Println(rt)

	return rt
}
